import taichi as ti
import random
import time
import numpy as np

ti.init(debug=True)

start_time = time.time()

n = 20

A = ti.field(dtype=ti.f32, shape=(n, n))
x = ti.field(dtype=ti.f32, shape=n)
new_x = ti.field(dtype=ti.f32, shape=n)
b = ti.field(dtype=ti.f32, shape=n)
D = ti.field(dtype=ti.f32, shape=(n, n))
I = ti.field(dtype=ti.f32, shape=(n, n))
B = ti.field(dtype=ti.f32, shape=(n, n))
f = ti.field(dtype=ti.f32, shape=n)

'''
对于Ax=b这个方程：
X_k+1 = Bj*X_k + fj
Bj = 
'''


@ti.kernel
def iterate():
    for row_B in range(n):
        t = 0.0
        for row_x in range(n):
            t+=x[row_x]*B[row_B,row_x]
        t+=f[row_B]
        new_x[row_B] = t
    for j in range(n):
        x[j]=new_x[j]
        
        # t = B@x
        # x = B@x+f


## 计算误差
@ti.kernel
def residual() -> ti.f32:
    res = 0.0
    # 误差向量e=b-A*x
    for i in range(n):
        r = b[i] * 1.0
        for j in range(n):
            r -= A[i, j] * x[j]
        res += r * r
    # res实际上是e的长度的平方    
    return res

for i in range(n):
    I[i,i]=1

for i in range(n):
    for j in range(n):
        A[i, j] = random.random() - 0.5
    A[i, i] += n * 0.1
    b[i] = random.random() * 100

for i in range(n):
    D[i,i] = 1/A[i,i]

np_B = I.to_numpy() - np.dot(D.to_numpy(),A.to_numpy())

np_f = np.dot(D.to_numpy(),b.to_numpy())

B.from_numpy(np_B)
f.from_numpy(np_f)


## 100次迭代
for i in range(100):
    iterate()
    print(f'iter {i}, residual={residual():0.10f}')
print(x)
    
end_time = time.time()
print("{}s".format(end_time-start_time))
